import React, { useState, useRef, useEffect, useCallback } from "react";
import {
  useReactTable,
  getCoreRowModel,
  getSortedRowModel,
  getFilteredRowModel,
  flexRender,
  ColumnDef,
  SortingState,
  RowSelectionState,
  createColumnHelper,
  Updater,
} from "@tanstack/react-table";
import {
  ChevronDown,
  ChevronUp,
  ExternalLink,
  Check,
  X,
  AlertTriangle,
  Shield,
  BarChart4,
  FileSearch,
  Eye,
} from "lucide-react";
import { cn } from "~/components/lib/utils";

// Define CveItem type for our table
export interface VulnTableItem {
  cve: string;
  severity: "CRITICAL" | "HIGH" | "MEDIUM" | "LOW" | "INFORMATIONAL";
  description: string;
  riskScore: number;
  affectedHosts: string[];
  published: string;
  lastModified?: string;
  references?: string[];
  cpe?: string[];
  fixed?: boolean;
  exploit?: boolean;
  selected?: boolean;
}

interface VulnerabilityCommandTableProps {
  data: VulnTableItem[];
  onSelectRow?: (selectedRows: string[]) => void;
  onViewDetails?: (cve: string) => void;
  onRefresh?: () => void;
  selectedRows?: string[];
  isLoading?: boolean;
}

export const VulnerabilityCommandTable = ({
  data,
  onSelectRow,
  onViewDetails,
  onRefresh,
  selectedRows = [],
  isLoading = false,
}: VulnerabilityCommandTableProps) => {
  // Table state
  const [sorting, setSorting] = useState<SortingState>([
    { id: "severity", desc: false },
  ]);
  const [rowSelection, setRowSelection] = useState<RowSelectionState>({});
  const [expandedRows, setExpandedRows] = useState<Record<string, boolean>>({});
  const [hoveredRowIndex, setHoveredRowIndex] = useState<string | null>(null);
  const [focusedRowIndex, setFocusedRowIndex] = useState<number | null>(null);
  const tableContainerRef = useRef<HTMLDivElement>(null);

  // Column definitions using column helper
  const columnHelper = createColumnHelper<VulnTableItem>();

  const columns = [
    // Selection column
    columnHelper.display({
      id: "select",
      header: ({ table }) => (
        <div className="px-1">
          <input
            type="checkbox"
            checked={table.getIsAllRowsSelected()}
            onChange={(e) => {
              table.getToggleAllRowsSelectedHandler()(e);
              e.stopPropagation(); // Prevent event bubbling
            }}
            onClick={(e) => {
              e.stopPropagation(); // Additional safeguard
              e.nativeEvent.stopImmediatePropagation(); // Full stop for event
            }}
            className="h-4 w-4 rounded border-gray-300 text-violet-600 focus:ring-violet-500 dark:border-gray-600 dark:bg-gray-700"
          />
        </div>
      ),
      cell: ({ row }) => (
        <div className="px-1" onClick={(e) => e.stopPropagation()}>
          <input
            type="checkbox"
            checked={row.getIsSelected()}
            onChange={(e) => {
              e.stopPropagation(); // Prevent row selection when clicking checkbox
              e.nativeEvent.stopImmediatePropagation(); // Full stop for event

              // Use setTimeout to ensure the event completes before we modify state
              setTimeout(() => {
                row.toggleSelected();
              }, 0);
            }}
            onClick={(e) => {
              e.stopPropagation(); // Prevent row click
              e.nativeEvent.stopImmediatePropagation(); // Additional safeguard
            }}
            className="h-4 w-4 rounded border-gray-300 text-violet-600 focus:ring-violet-500 dark:border-gray-600 dark:bg-gray-700"
          />
        </div>
      ),
      size: 40,
    }),

    // Severity column (added back)
    columnHelper.accessor("severity", {
      header: "Severity",
      cell: ({ getValue }) => {
        const severity = getValue();
        let bgColor = "";
        let textColor = "";

        switch (severity) {
          case "CRITICAL":
            bgColor = "bg-red-100 dark:bg-red-900/20";
            textColor = "text-red-800 dark:text-red-300";
            break;
          case "HIGH":
            bgColor = "bg-orange-100 dark:bg-orange-900/20";
            textColor = "text-orange-800 dark:text-orange-300";
            break;
          case "MEDIUM":
            bgColor = "bg-yellow-100 dark:bg-yellow-900/20";
            textColor = "text-yellow-800 dark:text-yellow-300";
            break;
          case "LOW":
            bgColor = "bg-green-100 dark:bg-green-900/20";
            textColor = "text-green-800 dark:text-green-300";
            break;
          case "INFORMATIONAL":
            bgColor = "bg-blue-100 dark:bg-blue-900/20";
            textColor = "text-blue-800 dark:text-blue-300";
            break;
        }

        return (
          <div
            className={`inline-flex rounded-full px-2 py-1 text-xs font-medium ${bgColor} ${textColor}`}
          >
            {severity.charAt(0) + severity.slice(1).toLowerCase()}
          </div>
        );
      },
      size: 100,
      enableSorting: true,
    }),

    // CVE ID column
    columnHelper.accessor("cve", {
      header: "CVE ID",
      cell: ({ getValue }) => (
        <div className="font-mono text-xs font-medium">{getValue()}</div>
      ),
      size: 140,
    }),

    // Risk Score column with visual gauge
    columnHelper.accessor("riskScore", {
      header: "CVSS",
      cell: ({ getValue }) => {
        const score = getValue();
        let color = "bg-blue-500";
        if (score >= 9.0) color = "bg-red-500";
        else if (score >= 7.0) color = "bg-orange-500";
        else if (score >= 4.0) color = "bg-yellow-500";
        else if (score > 0) color = "bg-green-500";

        return (
          <div className="flex items-center">
            <div className="mr-2 text-xs font-medium">{score.toFixed(1)}</div>
            <div className="h-1.5 w-12 rounded-full bg-gray-200 dark:bg-gray-700">
              <div
                className={`h-1.5 rounded-full ${color}`}
                style={{ width: `${Math.min(score * 10, 100)}%` }}
              ></div>
            </div>
          </div>
        );
      },
      size: 100,
      enableSorting: true,
    }),

    // Description column (expandable)
    columnHelper.accessor("description", {
      header: "Description",
      cell: ({ getValue, row }) => {
        const isExpanded = expandedRows[row.id];
        return (
          <div>
            <div
              className={cn(
                "text-xs text-gray-900 dark:text-gray-100",
                !isExpanded && "line-clamp-1"
              )}
            >
              {getValue()}
            </div>
            {getValue().length > 120 && (
              <button
                onClick={(e) => {
                  e.stopPropagation(); // Prevent row selection when clicking expand
                  toggleRowExpand(row.id);
                }}
                className="mt-0.5 text-xs text-violet-600 hover:text-violet-800 dark:text-violet-400 dark:hover:text-violet-300"
              >
                {isExpanded ? "Show less" : "Show more"}
              </button>
            )}
          </div>
        );
      },
      size: 450,
    }),

    // Actions column
    columnHelper.display({
      id: "actions",
      header: "",
      cell: ({ row }) => {
        const rowIndex = parseInt(row.id);
        const isHovered =
          hoveredRowIndex === row.id || focusedRowIndex === rowIndex;
        const cve = row.original.cve;

        return (
          <div
            className={cn(
              "flex items-center justify-end gap-1 pr-2 transition-opacity",
              isHovered ? "opacity-100" : "opacity-0 group-hover:opacity-100"
            )}
          >
            <button
              onClick={(e) => {
                e.stopPropagation();
                onViewDetails?.(cve);
              }}
              className="rounded p-1 text-gray-400 hover:bg-gray-100 hover:text-violet-600 dark:hover:bg-gray-700 dark:hover:text-violet-400"
              title="View vulnerability details"
            >
              <Eye className="h-3.5 w-3.5" />
            </button>
            <a
              href={`https://nvd.nist.gov/vuln/detail/${cve}`}
              target="_blank"
              rel="noopener noreferrer"
              className="rounded p-1 text-gray-400 hover:bg-gray-100 hover:text-blue-600 dark:hover:bg-gray-700 dark:hover:text-blue-400"
              title="View in NVD database"
              onClick={(e) => e.stopPropagation()}
            >
              <ExternalLink className="h-3.5 w-3.5" />
            </a>
          </div>
        );
      },
      size: 80,
    }),

    // Affected Hosts count (at the far right)
    columnHelper.accessor("affectedHosts", {
      header: "Hosts",
      cell: ({ getValue }) => {
        const hosts = getValue();
        return (
          <div className="text-center">
            <span className="text-xs font-medium text-gray-900 dark:text-gray-100">
              {hosts.length}
            </span>
          </div>
        );
      },
      size: 60,
      enableSorting: true,
    }),
  ] as ColumnDef<VulnTableItem>[];

  // Update the row selection effect to use a ref to prevent stale closures
  const prevSelectedRowsRef = useRef<string[]>(selectedRows);

  useEffect(() => {
    prevSelectedRowsRef.current = selectedRows;

    // Only update selection if the selected rows have actually changed
    const currentSelectedRowIds = Object.keys(rowSelection).map((idx) =>
      parseInt(idx)
    );
    const currentSelectedCVEs = currentSelectedRowIds
      .map((idx) => data[idx]?.cve)
      .filter(Boolean);

    // Check if the arrays are different
    const hasSelectionChanged =
      selectedRows.length !== currentSelectedCVEs.length ||
      selectedRows.some((cve) => !currentSelectedCVEs.includes(cve));

    if (hasSelectionChanged) {
      if (selectedRows.length === 0 && Object.keys(rowSelection).length > 0) {
        setRowSelection({});
      } else if (selectedRows.length > 0) {
        const newSelection: RowSelectionState = {};
        selectedRows.forEach((cve) => {
          const rowIndex = data.findIndex((item) => item.cve === cve);
          if (rowIndex !== -1) {
            newSelection[rowIndex] = true;
          }
        });
        setRowSelection(newSelection);
      }
    }
  }, [selectedRows, data, rowSelection]);

  // Improved callback for selection changes
  const handleSelectionChange = useCallback(
    (updaterOrValue: Updater<RowSelectionState>) => {
      // Handle updater function or direct value
      const getNewSelection = (old: RowSelectionState) => {
        return typeof updaterOrValue === "function"
          ? updaterOrValue(old)
          : updaterOrValue;
      };

      const newRowSelection = getNewSelection(rowSelection);
      setRowSelection(newRowSelection);

      if (onSelectRow) {
        const selectedCVEs = Object.keys(newRowSelection)
          .map((idx) => {
            const index = parseInt(idx);
            return index >= 0 && index < data.length
              ? data[index]?.cve
              : undefined;
          })
          .filter((cve): cve is string => cve !== undefined);

        onSelectRow(selectedCVEs);
      }
    },
    [data, onSelectRow, rowSelection]
  );

  // Set up the table with improved selection handling
  const table = useReactTable({
    data,
    columns,
    state: {
      sorting,
      rowSelection,
    },
    enableRowSelection: true,
    onRowSelectionChange: handleSelectionChange,
    onSortingChange: setSorting,
    getCoreRowModel: getCoreRowModel(),
    getSortedRowModel: getSortedRowModel(),
    getFilteredRowModel: getFilteredRowModel(),
  });

  // Toggle row expansion
  const toggleRowExpand = (rowId: string) => {
    setExpandedRows((prev) => ({
      ...prev,
      [rowId]: !prev[rowId],
    }));
  };

  // Completely separate row click from selection
  const handleRowClick = (rowIndex: number) => {
    setFocusedRowIndex(rowIndex);

    // Focus the container when a row is clicked to enable keyboard nav
    if (tableContainerRef.current) {
      tableContainerRef.current.focus();
    }
  };

  // Focus the container when the component mounts
  useEffect(() => {
    if (tableContainerRef.current && data.length > 0) {
      tableContainerRef.current.focus();
      // Set focus to the first row if no row is focused
      if (focusedRowIndex === null) {
        setFocusedRowIndex(0);
      }
    }
  }, [data, focusedRowIndex]);

  // Improve the space key handler to better manage selection
  const handleKeyDown = (e: React.KeyboardEvent) => {
    if (focusedRowIndex === null && data.length > 0) {
      setFocusedRowIndex(0);
      return;
    }

    if (focusedRowIndex === null) return;

    switch (e.key) {
      case "ArrowDown":
        e.preventDefault();
        if (focusedRowIndex < data.length - 1) {
          setFocusedRowIndex(focusedRowIndex + 1);
        }
        break;
      case "ArrowUp":
        e.preventDefault();
        if (focusedRowIndex > 0) {
          setFocusedRowIndex(focusedRowIndex - 1);
        }
        break;
      case "Enter":
        e.preventDefault();
        if (data[focusedRowIndex]) {
          onViewDetails?.(data[focusedRowIndex].cve);
        }
        break;
      case " ": // Spacebar
        e.preventDefault();
        e.stopPropagation();
        // Toggle selection of the focused row using the table's API
        if (focusedRowIndex !== null) {
          // Use setTimeout to defer state update until after event handling
          setTimeout(() => {
            const row = table.getRow(String(focusedRowIndex));
            const newSelectionState = !row.getIsSelected(); // Get the current state and invert it

            // Use the updater form to ensure we have the latest state
            table.setRowSelection((prev) => ({
              ...prev,
              [focusedRowIndex]: newSelectionState,
            }));
          }, 0);
        }
        break;
    }
  };

  return (
    <div
      ref={tableContainerRef}
      className="relative h-full overflow-auto outline-none" // added outline-none
      tabIndex={0} // Make the container focusable
      onKeyDown={handleKeyDown} // Handle keyboard events directly
    >
      <table className="w-full table-fixed border-collapse">
        <thead className="sticky top-0 z-10 bg-white dark:bg-gray-800">
          {table.getHeaderGroups().map((headerGroup) => (
            <tr
              key={headerGroup.id}
              className="border-b border-gray-200 dark:border-gray-700"
            >
              {headerGroup.headers.map((header) => (
                <th
                  key={header.id}
                  style={{ width: `${header.column.getSize()}px` }}
                  className={cn(
                    "bg-gray-50 px-3 py-2 text-left text-xs font-medium text-gray-500 dark:bg-gray-800 dark:text-gray-400",
                    header.column.getCanSort() && "cursor-pointer select-none"
                  )}
                  onClick={header.column.getToggleSortingHandler()}
                >
                  <div className="flex items-center">
                    {flexRender(
                      header.column.columnDef.header,
                      header.getContext()
                    )}
                    {header.column.getCanSort() && (
                      <div className="ml-1">
                        {header.column.getIsSorted() === "asc" ? (
                          <ChevronUp className="h-3 w-3" />
                        ) : header.column.getIsSorted() === "desc" ? (
                          <ChevronDown className="h-3 w-3" />
                        ) : (
                          <div className="h-3 w-3 opacity-0 group-hover:opacity-50">
                            ⋮
                          </div>
                        )}
                      </div>
                    )}
                  </div>
                </th>
              ))}
            </tr>
          ))}
        </thead>
        <tbody>
          {isLoading ? (
            <tr>
              <td colSpan={columns.length} className="py-10 text-center">
                <div className="flex flex-col items-center justify-center">
                  <div className="h-8 w-8 animate-spin rounded-full border-2 border-gray-300 border-t-violet-600"></div>
                  <span className="mt-2 text-sm text-gray-500">
                    Loading vulnerabilities...
                  </span>
                </div>
              </td>
            </tr>
          ) : table.getRowModel().rows.length === 0 ? (
            <tr>
              <td colSpan={columns.length} className="py-10 text-center">
                <div className="flex flex-col items-center justify-center">
                  <Shield className="h-10 w-10 text-gray-300 dark:text-gray-600" />
                  <span className="mt-2 text-sm text-gray-500">
                    No vulnerabilities found
                  </span>
                </div>
              </td>
            </tr>
          ) : (
            table.getRowModel().rows.map((row) => {
              const rowIndex = parseInt(row.id);
              return (
                <tr
                  key={row.id}
                  className={cn(
                    "group border-b border-gray-100 dark:border-gray-700/50",
                    hoveredRowIndex === row.id &&
                      "bg-gray-100 dark:bg-gray-700/50", // More visible hover state
                    row.getIsSelected() && "bg-violet-50 dark:bg-violet-900/20",
                    focusedRowIndex === rowIndex &&
                      "bg-blue-100 dark:bg-blue-900/30", // More visible focus state
                    "cursor-pointer transition-colors duration-100" // Added cursor-pointer
                  )}
                  onMouseEnter={() => setHoveredRowIndex(row.id)}
                  onMouseLeave={() => setHoveredRowIndex(null)}
                  onClick={(e) => {
                    // Identify if we're clicking a control
                    const target = e.target as Element;
                    const isClickingControl =
                      e.target instanceof HTMLInputElement ||
                      e.target instanceof HTMLButtonElement ||
                      e.target instanceof HTMLAnchorElement ||
                      (e.target instanceof SVGElement &&
                        e.target.closest("button, a")) ||
                      // Check parent elements up to 3 levels deep
                      (target.closest && target.closest("input")) ||
                      (target.closest && target.closest("button")) ||
                      (target.closest && target.closest("a"));

                    // Only proceed if not clicking a control
                    if (!isClickingControl) {
                      handleRowClick(rowIndex);

                      // Only toggle selection if not clicking a checkbox or button
                      if (!e.defaultPrevented) {
                        row.toggleSelected();
                      }
                    }
                  }}
                >
                  {row.getVisibleCells().map((cell) => (
                    <td key={cell.id} className="px-3 py-2 text-xs">
                      {flexRender(
                        cell.column.columnDef.cell,
                        cell.getContext()
                      )}
                    </td>
                  ))}
                </tr>
              );
            })
          )}
        </tbody>
      </table>

      {/* Keyboard shortcuts help tooltip */}
      <div className="absolute bottom-2 right-2 text-xs text-gray-400">
        <kbd className="rounded border border-gray-300 bg-gray-100 px-1 py-0.5 font-mono dark:border-gray-600 dark:bg-gray-800">
          ↑↓
        </kbd>{" "}
        Navigate
        <kbd className="ml-1 rounded border border-gray-300 bg-gray-100 px-1 py-0.5 font-mono dark:border-gray-600 dark:bg-gray-800">
          Space
        </kbd>{" "}
        Select
        <kbd className="ml-1 rounded border border-gray-300 bg-gray-100 px-1 py-0.5 font-mono dark:border-gray-600 dark:bg-gray-800">
          Enter
        </kbd>{" "}
        View Details
      </div>
    </div>
  );
};

export default VulnerabilityCommandTable;
